{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "import os\n",
    "import json\n",
    "import glob\n",
    "\n",
    "# 配置路径\n",
    "json_folder = '/root/PaddleOCR/dataset/en_rec_3W_res'      # 替换为你的 JSON 文件夹路径\n",
    "output_folder = './output_folder'  # 替换为输出 txt 文件夹路径\n",
    "image_prefix = 'images'                       # 图片路径前缀，如 \"images\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "正在搜索所有 JSON 文件...\n",
      "共找到 0 个 JSON 文件。\n",
      "训练集: 0 个文件\n",
      "验证集: 0 个文件\n",
      "已写入: /root/PaddleOCR/dataset/en_rec_3W_res/train.txt\n",
      "已写入: /root/PaddleOCR/dataset/en_rec_3W_res/val.txt\n",
      "✅ 所有文件处理完成！\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "import random\n",
    "from glob import glob\n",
    "\n",
    "# ================== 配置参数 ==================\n",
    "json_root_folder = '/root/PaddleOCR/dataset/en_rec_3W_res'      # 替换为你的根文件夹路径\n",
    "output_dir = '/root/PaddleOCR/dataset/en_rec_3W_res'          # 输出文件夹\n",
    "image_prefix = '1'                            # 图片路径前缀，如 images/\n",
    "\n",
    "# 创建输出目录\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "# 输出文件路径\n",
    "train_txt = os.path.join(output_dir, 'train.txt')\n",
    "val_txt = os.path.join(output_dir, 'val.txt')\n",
    "\n",
    "# ================== 1. 递归查找所有 .json 文件 ==================\n",
    "print(\"正在搜索所有 JSON 文件...\")\n",
    "json_files = []\n",
    "for root, dirs, files in os.walk(json_root_folder):\n",
    "    for file in files:\n",
    "        if file.lower().endswith('.json'):\n",
    "            json_files.append(os.path.join(root, file))\n",
    "\n",
    "print(f\"共找到 {len(json_files)} 个 JSON 文件。\")\n",
    "\n",
    "# ================== 2. 打乱顺序并划分 9:1 ==================\n",
    "random.seed(42)  # 固定随机种子，保证可复现\n",
    "random.shuffle(json_files)\n",
    "\n",
    "split_idx = int(0.9 * len(json_files))\n",
    "train_files = json_files[:split_idx]\n",
    "val_files = json_files[split_idx:]\n",
    "\n",
    "print(f\"训练集: {len(train_files)} 个文件\")\n",
    "print(f\"验证集: {len(val_files)} 个文件\")\n",
    "\n",
    "# ================== 3. 处理函数 ==================\n",
    "def process_json_file(json_path):\n",
    "    try:\n",
    "        with open(json_path, 'r', encoding='utf-8') as f:\n",
    "            data = json.load(f)\n",
    "    except Exception as e:\n",
    "        print(f\"读取失败: {json_path}, 错误: {e}\")\n",
    "        return None\n",
    "\n",
    "    # 获取 imagePath，若不存在则用文件名\n",
    "    image_filename = data.get(\"imagePath\")\n",
    "    if not image_filename:\n",
    "        # 使用 JSON 文件同名的图片\n",
    "        base_name = os.path.splitext(os.path.basename(json_path))[0]\n",
    "        image_filename = f\"{base_name}.png\"  # 可根据实际图片格式调整\n",
    "\n",
    "    print(image_filename)\n",
    "    image_path = f\"{image_prefix}/{image_filename}\"\n",
    "\n",
    "    # 提取 shapes\n",
    "    entries = []\n",
    "    for shape in data.get(\"shapes\", []):\n",
    "        desc = str(shape.get(\"description\", \"\")).strip()  # 转字符串并去空\n",
    "        transcription = desc if desc else \"###\"\n",
    "\n",
    "        # 坐标取整\n",
    "        points = [[int(round(float(coord))) for coord in point] for point in shape[\"points\"]]\n",
    "\n",
    "        entries.append({\n",
    "            \"transcription\": transcription,\n",
    "            \"points\": points\n",
    "        })\n",
    "\n",
    "    # 格式化为一行字符串\n",
    "    line = image_path + \"\\t\" + json.dumps(entries, ensure_ascii=False, separators=(',', ':'))\n",
    "    return line\n",
    "\n",
    "# ================== 4. 写入输出文件 ==================\n",
    "def write_lines(file_list, output_path):\n",
    "    with open(output_path, 'w', encoding='utf-8') as f:\n",
    "        for json_path in file_list:\n",
    "            line = process_json_file(json_path)\n",
    "            if line:\n",
    "                f.write(line + '\\n')\n",
    "    print(f\"已写入: {output_path}\")\n",
    "\n",
    "# 处理并写入\n",
    "write_lines(train_files, train_txt)\n",
    "write_lines(val_files, val_txt)\n",
    "\n",
    "print(\"✅ 所有文件处理完成！\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "训练集: 16542 个文件\n",
      "验证集: 1838 个文件\n",
      "\n",
      "成功提取 18380 个图片信息。\n",
      "结果已保存到: /root/PaddleOCR/dataset/en_rec_3w_res/train.txt\n",
      "\n",
      "成功提取 {'5': 6129, '7': 4293, '6': 8545, '9': 1609, 'a': 7840, 'r': 5784, 'e': 10297, 'c': 4038, 'l': 4068, 'd': 4530, '0': 1664, 'p': 2211, 'n': 7096, 'i': 8839, 'g': 2909, 'B': 174, 'w': 1846, '8': 1639, 'f': 1346, 'o': 4754, 'm': 1888, '1': 1753, 'b': 2001, 'u': 3000, 's': 6376, '2': 2369, 't': 7149, 'y': 1508, 'D': 114, 'C': 371, '3': 1826, '4': 3097, 'h': 3977, 'A': 299, 'k': 1157, 'S': 373, 'W': 327, 'P': 92, 'q': 372, 'H': 11, 'v': 1237, 'T': 111, 'M': 11, 'L': 45, 'I': 108, 'Z': 1, 'F': 15, 'O': 17, 'X': 1, 'x': 28, 'V': 27, 'N': 5, 'Q': 5, 'G': 130, 'E': 152, 'j': 11, 'U': 11, 'z': 39, 'K': 7, 'R': 3, 'J': 5, 'Y': 6}\n"
     ]
    }
   ],
   "source": [
    "import os,shutil,random\n",
    "from pathlib import Path\n",
    "import re\n",
    "\n",
    "\n",
    "def answer_to_pre(lable_answer):\n",
    "    \n",
    "    lable_answer=lable_answer.replace(\"ā\",\"a\").replace(\"ρ\",\"p\").replace(\"`\",\"\").replace(\"α\",\"p\") .replace(\"(\",\"\") .replace(\"`\",\"p\")  \n",
    "    # return lable_answer\n",
    "    \"\"\"只保留英文字母和数字\"\"\"\n",
    "    return re.sub(r'[^a-zA-Z0-9]', '', lable_answer)\n",
    "\n",
    "def extract_image_info_and_save(root_dir: str,DIRECTORY_OUTPUT_DIR: str, output_file: str = \"extracted_images.txt\",output_file_dict: str = \"extracted_images_dict.txt\"):\n",
    "    \"\"\"\n",
    "    遍历指定目录及其子目录，提取符合特定模式的图片文件名信息，并保存到文本文件中。\n",
    "    \n",
    "    假设图片文件名格式为: <label> <description>_<random1>_<random2>.<ext>\n",
    "    例如: 56 awarded_7076_2078379.png\n",
    "    提取后格式: relative_path_to_image label description\n",
    "    例如: images/train_word_1.png 56 awarded\n",
    "\n",
    "    Args:\n",
    "        root_dir (str): 要遍历的根目录路径\n",
    "        output_file (str): 输出的文本文件名，默认为 'extracted_images.txt'\n",
    "    \"\"\"\n",
    "\n",
    "    \n",
    "    # 定义支持的图片扩展名\n",
    "    image_extensions = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp'}\n",
    "\n",
    "    extracted_data = []\n",
    "    extracted_data_dict = {}\n",
    "    root_path = Path(root_dir)\n",
    "    \n",
    "    if not root_path.exists():\n",
    "        raise FileNotFoundError(f\"指定的目录不存在: {root_dir}\")\n",
    "    \n",
    "    if not root_path.is_dir():\n",
    "        raise NotADirectoryError(f\"指定的路径不是一个目录: {root_dir}\")\n",
    "    \n",
    "    if len(DIRECTORY_OUTPUT_DIR.strip())!=0:\n",
    "        if not os.path.isdir(DIRECTORY_OUTPUT_DIR):\n",
    "            os.makedirs(DIRECTORY_OUTPUT_DIR)\n",
    "        print(f\"处理后保存的路径:{DIRECTORY_OUTPUT_DIR}\")\n",
    "\n",
    "    # 遍历目录及其子目录\n",
    "    for file_path in root_path.rglob('*'):\n",
    "        if file_path.is_file() and file_path.suffix.lower() in image_extensions:\n",
    "\n",
    "            relative_path = file_path.relative_to(root_path)\n",
    "            # 使用正斜杠作为路径分隔符以确保跨平台兼容性\n",
    "            formatted_path = str(relative_path).replace('\\\\', '/')\n",
    "            if \" \" in formatted_path:\n",
    "                continue\n",
    "\n",
    "\n",
    "            # 提取标签和描述部分（第一组）\n",
    "            file_name = formatted_path.split(\"/\")[-1]\n",
    "            label_description=file_name.split(\"_\")[0].replace(' ', '')\n",
    "\n",
    "            #清除特殊字符\n",
    "            label_description=answer_to_pre(label_description)\n",
    "            \n",
    "            move_file_name=label_description+\"_\"+\"_\".join(file_name.split(\"_\")[-2:])\n",
    "\n",
    "            for i_ in label_description:\n",
    "                i=str(i_)\n",
    "                if i not in extracted_data_dict:\n",
    "                    extracted_data_dict[i]=0\n",
    "                extracted_data_dict[i]+=1\n",
    "                # print(i)\n",
    "\n",
    "            #图片重命名\n",
    "            if len(DIRECTORY_OUTPUT_DIR.strip())!=0:\n",
    "                new_dir_path=os.path.join(DIRECTORY_OUTPUT_DIR,\"/\".join(formatted_path.split(\"/\")[:-1]) )\n",
    "                if not os.path.isdir(new_dir_path):\n",
    "                    os.makedirs(new_dir_path)\n",
    "                base_file_path=os.path.join(root_path,formatted_path)\n",
    "                new_file_path=os.path.join(DIRECTORY_OUTPUT_DIR,formatted_path.replace(file_name,move_file_name))\n",
    "\n",
    "                #更新相对路径\n",
    "                # formatted_path_new = \"\".join(formatted_path.split(\"/\")[:-1]+[move_file_name])\n",
    "                formatted_path=formatted_path.replace(file_name,move_file_name)\n",
    "                # print(\"移动文件路径\",base_file_path,new_file_path)\n",
    "                shutil.copy(base_file_path,new_file_path)\n",
    "\n",
    "            # 获取相对于根目录的路径\n",
    "            try:\n",
    "                # 构建输出行\n",
    "                output_line = f\"{formatted_path} {label_description}\"\n",
    "                extracted_data.append(output_line)\n",
    "                \n",
    "                # print(f\"已处理: {output_line}\")\n",
    "                \n",
    "            except ValueError:\n",
    "                # 如果文件不在根目录下（理论上不会发生，因为使用了rglob）\n",
    "                print(f\"警告: 文件不在根目录内，跳过: {file_path}\")\n",
    "    \n",
    "    # 将结果写入文件\n",
    "    try:\n",
    "        # ================== 2. 打乱顺序并划分 9:1 ==================\n",
    "        random.seed(42)  \n",
    "        random.shuffle(extracted_data)\n",
    "\n",
    "        split_idx = int(0.9 * len(extracted_data))\n",
    "        train_files = extracted_data[:split_idx]\n",
    "        val_files = extracted_data[split_idx:]\n",
    "\n",
    "        print(f\"训练集: {len(train_files)} 个文件\")\n",
    "        print(f\"验证集: {len(val_files)} 个文件\")\n",
    "\n",
    "\n",
    "        with open(output_file, 'w', encoding='utf-8') as f:\n",
    "            for line in train_files:\n",
    "                f.write(line + '\\n')\n",
    "\n",
    "        with open(output_file.replace(\"train.txt\",\"val.txt\"), 'w', encoding='utf-8') as f:\n",
    "            for line in val_files:\n",
    "                f.write(line + '\\n')\n",
    "        print(f\"\\n成功提取 {len(extracted_data)} 个图片信息。\")\n",
    "        print(f\"结果已保存到: {os.path.abspath(output_file)}\")\n",
    "        \n",
    "        with open(output_file_dict, 'w', encoding='utf-8') as f:\n",
    "            for line in extracted_data_dict:\n",
    "                f.write(line + '\\n')\n",
    "        print(f\"\\n成功提取 {extracted_data_dict}\")\n",
    "        \n",
    "    except IOError as e:\n",
    "        print(f\"写入文件时发生错误: {e}\")\n",
    "        raise\n",
    "\n",
    "# -------------------------- 使用示例 --------------------------\n",
    "    \n",
    "if __name__ == \"__main__\":\n",
    "    # 请修改这里的路径为您的实际目录\n",
    "    DIRECTORY_TO_SCAN = \"/root/PaddleOCR/dataset/en_rec_3w_res\"  # <-- 修改这里\n",
    "    DIRECTORY_OUTPUT_DIR = \"\"  # <-- 需要修改文件命名配置此路径，不需要为空即可\n",
    "    OUTPUT_FILENAME = \"/root/PaddleOCR/dataset/en_rec_3w_res/train.txt\"      # <-- 可选：修改输出文件名\n",
    "    OUTPUT_FILENAME_dict = \"/root/PaddleOCR/dataset/en_rec_3w_res/en_dict.txt\"      # <-- 可选：修改输出文件名\n",
    "    \n",
    "    try:\n",
    "        extract_image_info_and_save(DIRECTORY_TO_SCAN,DIRECTORY_OUTPUT_DIR, OUTPUT_FILENAME,OUTPUT_FILENAME_dict)\n",
    "    except (FileNotFoundError, NotADirectoryError) as e:\n",
    "        print(f\"错误: {e}\")\n",
    "        print(\"请检查您设置的目录路径是否正确。\")\n",
    "    except Exception as e:\n",
    "        print(f\"发生未预期的错误: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
